# -*- coding: utf-8 -*-
# 对所有样本依次计算时频图 并保存
import matplotlib.pyplot as plt
import numpy as np
import pywt
from scipy.io import loadmat
import os
import shutil # 用于与操作系统进行交互，如创建文件夹、删除文件夹
import scipy.io as sio
import mne
shutil.rmtree(r'D:\Project_mb\CWTImages\BCI42a') # 会删除名为'image/CWT/BCI_IV2a'的目录及其所有内容。rmtree函数用于删除整个目录树
os.mkdir(r'D:\Project_mb\CWTImages\BCI42a') # 创建一个名为'iimage/CWT/BCI_IV2a'的新目录。
# 生成CWT时频图像
def makeTimeFrequencyImage(data, img_path, img_size,fs,sfreq, totalscales,wavename='cmor3-3'):
    """
    将一维脑电信号转为二维时频图像
    :param data: 一维时间序列
    :param img_path: 图像存储的路径
    :param img_size: 图像存储的尺寸
    :param sampling_rate: 数据的采样率
    :param totalscal: 尺度长度
    :param wavename: 小波基函数
    :return:
    """
    data = np.array(data)
    sampling_length = len(data)
    # 尺度对应的频率范围
    scales = np.arange(1, totalscales)
    frequencies = pywt.scale2frequency(wavename, scales) * fs
    selected_scales = scales[(frequencies >= sfreq[0]) & (frequencies <= sfreq[1])]
    # 进行连续小波变换
    coefficients, frequencies = pywt.cwt(data, selected_scales, wavename, sampling_period=1/fs)
    # 系数矩阵绝对值
    amp = abs(coefficients)
    # 生成图片
    t = np.linspace(0, sampling_length/fs, sampling_length, endpoint=False)
    # plt.contourf(t, frequencies, amp, cmap='jet')
    plt.contourf(t, frequencies, amp, cmap='jet')
    plt.axis('off')  # 设置图像坐标轴不可见
    plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)  # 调整子图的位置和间距，将子图充满整个图像
    plt.margins(0, 0)  # 设置图像的边距为0，即没有额外的空白边框。
    plt.gcf().set_size_inches(img_size / 100, img_size / 100)  # 设置图像的大小，单位为英寸
    plt.savefig(img_path, dpi=100)
    plt.clf()  # 避免内存溢出
    plt.close()  # 释放内存
def GenerateImageDataset(dataPath, savepath, img_size,sfreq, totalscales=256,wavename='cmor3-3', session=1):
    data_dict = sio.loadmat(dataPath)
    data = data_dict['data']
    labels = data_dict['labels'].reshape(-1)
    channel_names = data_dict['channel_names']
    fs = int(data_dict['fs'])
    for i in range(data.shape[0]):  # trials:288
        for j in range(data.shape[1]):  # channels:22
            channel_label = str(channel_names[j]).strip()
            # # if channel_label == 'C3':
            # channels = ['C1','C2','C3','C4','C5','C6','Cz']
            # if channel_label in channels:
            # if channel_label == 'C3' or channel_label == 'C4' or channel_label == 'Cz':
            single_data = data[i, j, :]
            label = labels[i]
            img_path = os.path.join(savepath, str(label)+"_"+str(i)+'_'+str(channel_label)+"_"+str(session))
            makeTimeFrequencyImage(single_data, img_path, img_size=img_size,fs=fs, sfreq=sfreq, totalscales=totalscales, wavename=wavename)



def BCI42a_image(dataPath, savePath, img_size, sfreq, totalscales,wavename='cmor3-3'):
    # 遍历目录下的所有文件和子目录
    for root, dirs, files in os.walk(dataPath):
        # root 表示当前目录路径
        # dirs 表示当前目录下的子目录名列表
        # files 表示当前目录下的文件名列表
        for file in files:
            data_savePath = savePath + "\\"+'sub' + str(file[1:3])
            # 检查文件是否存在
            if not os.path.exists(data_savePath):
                os.mkdir(data_savePath)
            if file[3] == 'T':
                session = 1
            #     data_save_path = data_savePath + '\\' + 'train'
            #     os.mkdir(data_save_path)
            elif file[3] == 'E':
                session = 2
            #     data_save_path = data_savePath + '\\' + 'test'
            #     os.mkdir(data_save_path)

            datPath = os.path.join(root, file)
            GenerateImageDataset(datPath, data_savePath,img_size=img_size, sfreq=sfreq,totalscales=totalscales, wavename=wavename, session = session)
            print(file+' Done!')




BCI42a_image(r'D:\Project_mb\data\BCI42a\preprocession_matdata', r'D:\Project_mb\CWTImages\BCI42a', img_size=224,sfreq=[4,40], totalscales=224, wavename='cmor3-3')

